package cn._51doit.flink.day06;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.StateTtlConfig;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

/**
 * 设置MapState的TTL
 * 比如按照省份进行KeyBy，然后设置MapState的TTL，那么是将每个小key即城市设置TTL
 */
public class MapStateTTLDemo {


    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        //开启checkpoint
        env.enableCheckpointing(10000);

        //辽宁省,沈阳市,2000
        //山东省,济南市,3000
        //辽宁省,大连市,1000
        //辽宁省,鞍山市,1000
        //山东省,烟台市,1000
        DataStreamSource<String> lines = env.socketTextStream("localhost", 8888);

        SingleOutputStreamOperator<Tuple3<String, String, Integer>> tpStream = lines.map(new MapFunction<String, Tuple3<String, String, Integer>>() {
            @Override
            public Tuple3<String, String, Integer> map(String s) throws Exception {
                String[] fields = s.split(",");
                String province = fields[0];
                String city = fields[1];
                int money = Integer.parseInt(fields[2]);
                return Tuple3.of(province, city, money);
            }
        });


        //相同省份的数据分到到一个分区内，并且累加每个城市的金额
        KeyedStream<Tuple3<String, String, Integer>, String> keyed = tpStream.keyBy(t -> t.f0);

        SingleOutputStreamOperator<Tuple3<String, String, Integer>> res = keyed.map(new RichMapFunction<Tuple3<String, String, Integer>, Tuple3<String, String, Integer>>() {

            private MapState<String, Integer> mapState;

            @Override
            public void open(Configuration parameters) throws Exception {
                //定义状态描述器
                StateTtlConfig ttlConfig = StateTtlConfig.newBuilder(Time.seconds(30)).build();
                MapStateDescriptor<String, Integer> stateDescriptor = new MapStateDescriptor<>("income-state", String.class, Integer.class);
                stateDescriptor.enableTimeToLive(ttlConfig);
                mapState = getRuntimeContext().getMapState(stateDescriptor);
            }

            @Override
            public Tuple3<String, String, Integer> map(Tuple3<String, String, Integer> tp) throws Exception {
                String city = tp.f1;
                Integer money = tp.f2;
                Integer historyMoney = mapState.get(city);
                if (historyMoney == null) {
                    historyMoney = 0;
                }
                historyMoney += money;
                //更新状态
                mapState.put(city, historyMoney);
                //输出数据
                tp.f2 = historyMoney;
                return tp;
            }
        });

        res.print();


        env.execute();
    }

}
